use std::{iter::Peekable, slice::Iter};

use nand2tetris_vmtranslator_rust::data::{ArithmeticCommand, Segment};

use crate::{
    compiler::ProgramCompilerUnit,
    data::{KeywordType, SymbolType},
    error::{CompilerError, CompilerResult},
    token::Token,
};

/// Statement compiler unit
pub(crate) struct StatementCompilerUnit;

impl StatementCompilerUnit {
    /// The compiler flow unit of statements
    pub(super) fn compile_statements(
        compiler: &mut ProgramCompilerUnit,
        tokens: &mut Peekable<Iter<'_, Token>>,
    ) -> CompilerResult<()> {
        while let Some(token) = tokens.peek()
            && **token != Token::Symbol(SymbolType::CloseBrace)
        {
            Self::compile_statement(compiler, tokens)?;
        }

        Ok(())
    }
    /// The compiler flow unit of statement
    fn compile_statement(
        compiler: &mut ProgramCompilerUnit,
        tokens: &mut Peekable<Iter<'_, Token>>,
    ) -> CompilerResult<()> {
        let token = tokens
            .peek()
            .ok_or_else(|| CompilerError::Statement("not token".into()))?;
        match token {
            Token::Keyword(KeywordType::Let) => Self::compile_let_statement(compiler, tokens),
            Token::Keyword(KeywordType::If) => Self::compile_if_statement(compiler, tokens),
            Token::Keyword(KeywordType::While) => Self::compile_while_statement(compiler, tokens),
            Token::Keyword(KeywordType::Do) => Self::compile_do_statement(compiler, tokens),
            Token::Keyword(KeywordType::Return) => Self::compile_return_statement(compiler, tokens),
            _ => Err(CompilerError::Statement(
                "not a let/if/while/do/return statement".into(),
            )),
        }
    }
    /// The compiler flow unit of return statement
    fn compile_return_statement(
        compiler: &mut ProgramCompilerUnit,
        tokens: &mut Peekable<Iter<'_, Token>>,
    ) -> CompilerResult<()> {
        // return
        let _ = crate::analyzer::Analyzer::match_keyword(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Statement("must start with return".into()))?,
            KeywordType::Return,
        )?;
        let mut is_return_void = false;
        if let Some(Token::Symbol(SymbolType::Semicolon)) = tokens.peek() {
            let _ = tokens.next();
            is_return_void = true;
        } else {
            // expression
            super::expression::ExpressionCompilerUnit::compile_expression(compiler, tokens)?;
            // ;
            let _ = crate::analyzer::Analyzer::match_symbol(
                tokens
                    .next()
                    .ok_or_else(|| CompilerError::Statement("miss ;".into()))?,
                SymbolType::Semicolon,
            )?;
        }
        // write return code
        if is_return_void {
            // void return 0
            compiler.vm_writer.write_push(Segment::Constant, 0);
        }
        compiler.vm_writer.write_return();
        Ok(())
    }

    /// The compiler flow unit of do statement
    fn compile_do_statement(
        compiler: &mut ProgramCompilerUnit,
        tokens: &mut Peekable<Iter<'_, Token>>,
    ) -> CompilerResult<()> {
        // do
        let _ = crate::analyzer::Analyzer::match_keyword(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Statement("must start with do".into()))?,
            KeywordType::Do,
        )?;
        // subroutineCall
        super::expression::ExpressionCompilerUnit::compile_subroutine_call(compiler, tokens, None)?;

        // ;
        let _ = crate::analyzer::Analyzer::match_symbol(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Statement("miss ;".into()))?,
            SymbolType::Semicolon,
        )?;
        // yark the return value
        compiler.vm_writer.write_pop(Segment::Temp, 0);

        Ok(())
    }

    /// The compiler flow unit of while statement
    fn compile_while_statement(
        compiler: &mut ProgramCompilerUnit,
        tokens: &mut Peekable<Iter<'_, Token>>,
    ) -> CompilerResult<()> {
        // while
        let _ = crate::analyzer::Analyzer::match_keyword(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Statement("must start with while".into()))?,
            KeywordType::While,
        )?;
        // (
        let _ = crate::analyzer::Analyzer::match_symbol(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Statement("miss (".into()))?,
            SymbolType::OpenParenthesis,
        )?;
        let true_label = compiler.get_label_counter();
        let false_label = compiler.get_label_counter();
        compiler.vm_writer.write_label(&true_label);
        // expression
        super::expression::ExpressionCompilerUnit::compile_expression(compiler, tokens)?;
        // )
        let _ = crate::analyzer::Analyzer::match_symbol(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Statement("miss )".into()))?,
            SymbolType::CloseParenthesis,
        )?;
        // {
        let _ = crate::analyzer::Analyzer::match_symbol(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Statement("miss {".into()))?,
            SymbolType::OpenBrace,
        )?;
        compiler
            .vm_writer
            .write_arithmetic_command(ArithmeticCommand::Not);
        compiler.vm_writer.write_if_goto(&false_label);
        // statements
        Self::compile_statements(compiler, tokens)?;
        // }
        let _ = crate::analyzer::Analyzer::match_symbol(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Statement("miss }".into()))?,
            SymbolType::CloseBrace,
        )?;

        compiler.vm_writer.write_goto(&true_label);
        compiler.vm_writer.write_label(&false_label);
        Ok(())
    }

    /// The compiler flow unit of if statement
    fn compile_if_statement(
        compiler: &mut ProgramCompilerUnit,
        tokens: &mut Peekable<Iter<'_, Token>>,
    ) -> CompilerResult<()> {
        // if
        let _ = crate::analyzer::Analyzer::match_keyword(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Statement("must start with if".into()))?,
            KeywordType::If,
        )?;
        // (
        let _ = crate::analyzer::Analyzer::match_symbol(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Statement("miss (".into()))?,
            SymbolType::OpenParenthesis,
        )?;
        // expression
        super::expression::ExpressionCompilerUnit::compile_expression(compiler, tokens)?;
        // )
        let _ = crate::analyzer::Analyzer::match_symbol(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Statement("miss )".into()))?,
            SymbolType::CloseParenthesis,
        )?;
        compiler
            .vm_writer
            .write_arithmetic_command(ArithmeticCommand::Not);

        let false_label = compiler.get_label_counter();
        let true_label = compiler.get_label_counter();

        compiler.vm_writer.write_if_goto(&false_label);

        // {
        let _ = crate::analyzer::Analyzer::match_symbol(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Statement("miss {".into()))?,
            SymbolType::OpenBrace,
        )?;
        // statements
        Self::compile_statements(compiler, tokens)?;
        // }
        let _ = crate::analyzer::Analyzer::match_symbol(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Statement("miss }".into()))?,
            SymbolType::CloseBrace,
        )?;
        compiler.vm_writer.write_goto(&true_label);

        compiler.vm_writer.write_label(&false_label);
        if let Some(Token::Keyword(KeywordType::Else)) = tokens.peek() {
            // skip [
            let _ = tokens.next();
            // {
            let _ = crate::analyzer::Analyzer::match_symbol(
                tokens
                    .next()
                    .ok_or_else(|| CompilerError::Statement("miss {".into()))?,
                SymbolType::OpenBrace,
            )?;
            // statements
            Self::compile_statements(compiler, tokens)?;
            // }
            let _ = crate::analyzer::Analyzer::match_symbol(
                tokens
                    .next()
                    .ok_or_else(|| CompilerError::Statement("miss }".into()))?,
                SymbolType::CloseBrace,
            )?;
        }
        compiler.vm_writer.write_label(&true_label);
        Ok(())
    }
    /// The compiler flow unit of let statement
    fn compile_let_statement(
        compiler: &mut ProgramCompilerUnit,
        tokens: &mut Peekable<Iter<'_, Token>>,
    ) -> CompilerResult<()> {
        // let
        let _ = crate::analyzer::Analyzer::match_keyword(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Statement("must start with let".into()))?,
            KeywordType::Let,
        )?;
        // varname
        let var_name = super::expression::ExpressionCompilerUnit::flow_unit_identifier(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Statement("miss var name".into()))?,
        )?;
        let var_name = var_name.as_str();
        // query var info
        let var_info = compiler
            .get_variable(var_name)
            .map(|item| item.to_owned())
            .ok_or_else(|| {
                CompilerError::Indentifier(format!("the variable is not exist : {var_name}"))
            })?;
        let mut is_array = false;
        if let Some(Token::Symbol(SymbolType::OpenBracket)) = tokens.peek() {
            // skip [
            let _ = tokens.next();
            // step 1 : push arr
            compiler
                .vm_writer
                .write_push(var_info.var_kind.into(), var_info.index_of_kind);

            super::expression::ExpressionCompilerUnit::compile_expression(compiler, tokens)?;
            let _ = crate::analyzer::Analyzer::match_symbol(
                tokens
                    .next()
                    .ok_or_else(|| CompilerError::SubroutineCall("miss ]".into()))?,
                SymbolType::CloseBracket,
            )?;

            // save the addr ( = arr + expression ) to stack
            // write add
            compiler
                .vm_writer
                .write_arithmetic_command(ArithmeticCommand::Add);
            is_array = true;
        }

        let _ = crate::analyzer::Analyzer::match_symbol(
            tokens
                .next()
                .ok_or_else(|| CompilerError::SubroutineCall("miss =".into()))?,
            SymbolType::Equal,
        )?;
        super::expression::ExpressionCompilerUnit::compile_expression(compiler, tokens)?;
        // ;
        let _ = crate::analyzer::Analyzer::match_symbol(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Statement("miss ;".into()))?,
            SymbolType::Semicolon,
        )?;
        // output vm code
        if is_array {
            // stage 1 : the variable is a array.
            // write the top value to  the that
            // save value to temp
            compiler.vm_writer.write_pop(Segment::Temp, 0);
            // set that to array's address : that = arr+expression
            compiler.vm_writer.write_pop(Segment::Pointer, 1);
            compiler.vm_writer.write_push(Segment::Temp, 0);
            compiler.vm_writer.write_pop(Segment::That, 0);
        } else {
            // stage 2 : the variable is not a array.
            // pop stack top value to variable
            compiler
                .vm_writer
                .write_pop(var_info.var_kind.into(), var_info.index_of_kind);
        }
        Ok(())
    }
}
